"""
本脚本用于探索外交关系的agent的prompt
"""
#%% init project env
import EsyPro
import os
import json
import torch
from typing import Dict, Any
from VicAi.chat_pai import MultiTurnChat

sfs = EsyPro.ScriptResultManager('zqf', locals(), version=0)

#%% prompt definitions

def role_prompt() -> str:
    """生成角色描述提示"""
    return (
        "你是一个精通历史外交关系推理的外交关系评论家，你正在对一段模拟历史进行评论，你喜欢这种模拟。"
        "这种模拟按照周期进行，你喜欢在每一个周期的开始询问当前的世界基本信息，然后筛选值得评论的事件，"
        "发表评论并解释，最后预测外交关系的变化。你会等待我的指令，然后按照指令参与模拟。"
        "现在你是一个精通历史外交关系推理的外交关系评论家，拥有深厚的历史知识储备，熟悉古今中外各种外交事件的来龙去脉以及处理方式，"
        "对国际关系理论、外交政策制定、国际法等相关知识了如指掌。"
        "你擅长运用这些专业知识，在模拟历史的环境中，对各国之间的外交动态进行精准分析和评论。"
        "你清楚地知道，这种模拟是按照周期进行的，在每一个周期的开始，你会主动询问当前的世界基本信息，"
        "包括重大外交决策（外交行动、博弈）、当前的外交关系以及重大科技、军事、社会和经济事件等，"
        "这是你发挥专业能力进行评论的基础。你善于从这些信息中筛选出值得评论的重要事件，"
        "凭借你的专业知识和敏锐洞察力，发表独到的评论并详细解释，同时还能结合当前形势和历史经验，"
        "对外交关系的变化进行准确预测。现在周期并未开始，处于准备阶段，"
        "请你仅用一段简洁的话，复述你对自身身份的理解，不要涉及其他内容，也不要采取任何行动。"
    )

def save_results(content: Dict, res: Dict, base_path: str = "/data/zqf/T20240921_yoao/SeekAnno/_l_zqf_s1_test_v0") -> None:
    """保存查询结果到文件"""
    os.makedirs(base_path, exist_ok=True)
    
    query_path = os.path.join(base_path, "query_answer.json")
    with open(query_path, 'w', encoding='utf-8') as f:
        json.dump(content, f, ensure_ascii=False, indent=2)
        f.write('\n')
        json.dump(res, f, ensure_ascii=False, indent=2)

#%% main
if __name__ == '__main__':
    print(f'start {__file__}')
    
    # 初始化基础路径
    BASE_PATH = "/data/zqf/T20240921_yoao/SeekAnno/_l_zqf_s1_test_v0"
    
    try:
        # 加载上下文
        context_path = os.path.join(BASE_PATH, "query_context.pth")
        content = torch.load(context_path)
        
        # 初始化多轮对话类
        chat = MultiTurnChat(model="deepseek-r1:8b-40k")
        
        # 初始角色提示
        initial_prompt = role_prompt()
        chat.add_user_message(initial_prompt)
        print(initial_prompt)
        
        # 第一次模型调用
        try:
            reply = chat.get_reply()
            print(f"AI: {reply}")
        except RuntimeError as e:
            print(f"模型调用失败: {e}")
        
        # 保存结果
        save_results(content, {"initial_reply": reply}, BASE_PATH)
        
        # 第二次模型调用
        follow_up_prompt = f"现在开始第1个周期的基础信息获取阶段，请按照我的指令参与模拟。"
        chat.add_user_message(follow_up_prompt)
        try:
            follow_up_reply = chat.get_reply()
            print(f"AI: {follow_up_reply}")
        except RuntimeError as e:
            print(f"模型调用失败: {e}")
        
    except Exception as e:
        print(f"程序执行出错: {str(e)}")
    
    print(f'end {__file__}')
